<template>
  <a-layout class="common-layout">
    <a-layout-sider :width="250" class="common-sider">
      <MyMenu :menu="menu" v-model:select-key="selectKey"></MyMenu>
    </a-layout-sider>
    <a-layout class="common-box">
      <a-layout-content>
        <!--        <component :is="typeComponentMap[component]" ></component>-->
        <div v-show="component === 'Interation'">
          <div class="Interation-box">
            <div class="main">
              <canvas ref="drawCanvas"></canvas>
            </div>
            <div class="button-box">
              <a-space :size='16' direction="vertical">
                <a-color-picker defaultValue="#165DFF" hideTrigger showHistory showPreset @change="colorChange"
                  :historyColors="history" />
                <a-button @click="predict" long type="primary">开始预测</a-button>
                <a-button @click="clear" long type="primary">清除画板</a-button>
                <a-button long type="primary"> 退出</a-button>
                <a-button long type="primary"> 保存作品 </a-button>
                <a-space></a-space>
                <a-typography-text> 画笔大小</a-typography-text>
                <a-input-number :model-value="size" class="input-demo" :min="5" :max="15" />
              </a-space>

            </div>


          </div>
        </div>
        <div v-if="component === 'InputExample'">
          <div ref="tfvisContainer"></div>
        </div>
        <div v-else-if="component === 'Architecture'">
          <div ref="ArchitectureRef"></div>
        </div>
        <div v-show="show">
          <div ref="TrainingRef"></div>
        </div>
      </a-layout-content>
    </a-layout>
  </a-layout>
</template>


<script setup>


import { computed, onMounted, ref, watch } from "vue";
import MyMenu from "../components/MyMenu.vue";
import { Message } from "@arco-design/web-vue";
import * as tf from '@tensorflow/tfjs'
import { Modal } from "@arco-design/web-vue";
import { run, showExamples } from "../utils/script";
import { MnistData } from "../utils/data";
import * as tfvis from "@tensorflow/tfjs-vis";


// const typeComponentMap = {
//   "Interation": Interation,
//   "InputExample": InputExample
// }

const menu = [
  {
    key: "1",
    title: "Input Data",
    icon: "icon-file-audio",
    children: [
      {
        key: "1-1",
        title: "Input Examples",
        componentName: "InputExample"
      },
    ]
  },
  {
    key: "2",
    title: "Model",
    icon: "icon-stamp",
    children: [
      {
        key: "2-1",
        title: "Architecture",
        componentName: "Architecture"
      },
      {
        key: "2-2",
        title: "Training",
        componentName: "Training"
      },
    ]
  },
  // {
  //   key: "3",
  //   title: "Evaluation",
  //   icon: "icon-gift",
  //   children: [
  //     {
  //       key: "3-1",
  //       title: "Accuracy",
  //     },
  //     {
  //       key: "3-2",
  //       title: "Confusion Matrix",
  //     },
  //   ]
  // },
  {
    key: "4",
    title: "Interaction",
    icon: "icon-highlight",
    componentName: "Interation"
  }
]

const selectKey = ref('4')


const drawCanvas = ref(null)
// const preview = ref(null)
const cvsWidth = ref(380);

const color = ref('#FFFFFF')

const history = ref(['#FFFFFF'])

const colorChange = (value) => {
  color.value = value
  addHistory(false, value)
}
const addHistory = (visible, color) => {
  if (!visible) {
    const index = history.value.indexOf(color);
    if (index !== -1) {
      history.value.splice(index, 1);
    }
    history.value.unshift(color);
  }
}

const size = defineModel({
  default: 10
})
const tfvisContainer = ref(null);
const ArchitectureRef = ref(null)
const TrainingRef = ref(null)

let model
let data

const show = computed(() => {
  return component.value === 'Training'
})

const component = computed(() => {
  console.log(selectKey.value)
  for (let i = 0; i < menu.length; i++) {
    if (menu[i].children) {
      if (menu[i].children.find((item) => item.key === selectKey.value)) {
        return menu[i].children.find((item) => item.key === selectKey.value).componentName
      }
    }
    else {
      if (menu[i].key === selectKey.value) {
        console.log(menu[i].componentName)
        return menu[i].componentName
      }
    }
  }
})


onMounted(async () => {
  let flag = localStorage.getItem('modelFlag')

  if (flag) {
    Message.info('模型已经训练好，可直接使用')
    model = await tf.loadLayersModel('localstorage://my-model');
  }
  else {
    Message.info('未检测到模型，开始训练')
    model = await run(TrainingRef)
    localStorage.setItem('modelFlag', true)
    await model.save('localstorage://my-model');
  }
  console.log('run finish')
  const canvas = drawCanvas.value
  console.log(canvas)
  const ctx = canvas.getContext('2d');
  canvas.width = cvsWidth.value
  canvas.height = cvsWidth.value
  // 绑定鼠标事件：按住左键移动绘制线条(利用矩阵连起来书写数字)
  canvas.addEventListener("mousemove", (e) => {
    if (e.buttons === 1) {
      // ctx.fillStyle = 'rgb(255,255,255)'
      ctx.fillStyle = color.value
      console.log(color)
      //绘制圆
      ctx.beginPath()
      console.log(ctx)
      ctx.arc(e.offsetX, e.offsetY, size.value, 0, Math.PI * 2)
      ctx.fill()
      // ctx.fillRect(e.offsetX, e.offsetY, cvsWeight.value, cvsWeight.value)

    }
  })
  clear()
})


watch(() => tfvisContainer.value, async (value) => {
  if (value) {
    data = new MnistData();
    await data.load();
    await showExamples(data, tfvisContainer);
  }
})


watch(() => ArchitectureRef.value, async (value) => {
  if (value) {
    await tfvis.show.modelSummary(ArchitectureRef.value, model);
  }
});



const clear = () => {
  const ctx = drawCanvas.value.getContext('2d');
  ctx.fillStyle = 'rgb(0,0,0)', ctx.fillRect(0, 0, cvsWidth.value, cvsWidth.value)
}



const predict = () => {
  const cvs = drawCanvas.value
  const imgData = cvs.getContext('2d').getImageData(0, 0, cvsWidth.value, cvsWidth.value)
  const color = imgData.data
  //裁剪，找到x轴和y轴中的白色，然后裁剪
  let x1 = cvsWidth.value
  let x2 = 0
  let y1 = cvsWidth.value
  let y2 = 0
  for (let i = 0; i < color.length; i += 4) {
    if (color[i] !== 0) {
      let x = (i / 4) % cvsWidth.value
      let y = Math.floor((i / 4) / cvsWidth.value)
      x1 = Math.min(x1, x)
      x2 = Math.max(x2, x)
      y1 = Math.min(y1, y)
      y2 = Math.max(y2, y)
    }
  }
  let width = Math.max(x2 - x1, y2 - y1)
  let startX = x1 - (width - (x2 - x1)) / 2
  let startY = y1 - (width - (y2 - y1)) / 2
  const imgData2 = cvs.getContext('2d').getImageData(startX, startY, width, width)
  const color2 = imgData2.data
  for (let i = 0; i < color2.length; i += 4) {
    if (color2[i] !== 0) {
      color2[i] = 255
      color2[i+1] = 255
      color2[i+2] = 255
      color2[i+3] = 255
    }
  }
  const newCanvas = document.createElement('canvas')
  const newCtx = newCanvas.getContext('2d')
  newCanvas.width = width
  newCanvas.height = width


  newCtx.putImageData(imgData2, 0, 0)
  // if (preview.value.lastChild.tagName ==='CANVAS'){
  //   preview.value.removeChild(preview.value.lastChild)
  // }
  // preview.value.appendChild(newCanvas)
  // document.body.appendChild(newCanvas)
  // 将canvas转换成Tensor，形状是28*28，黑白图片,并归一化
  const input = tf.tidy(() => {
    return tf.image.resizeBilinear(
      tf.browser.fromPixels(newCanvas),
      [28, 28],
      true,
    )
      .slice([0, 0, 0], [28, 28, 1])
      .toFloat()
      .div(255)
      .reshape([1, 28, 28, 1])

  });
  // 预测
  const pred = model.predict(input).argMax(1);
  Modal.info({
    title: '预测结果',
    content: `预测结果为：${pred.dataSync()[0]}`

  })
  // alert(`预测结果为：${pred.dataSync()[0]}`)
  // clear()
}

</script>


<style scoped>
.common-layout {
  min-height: 100vh;
  min-width: 100vw;

  .common-sider {}

  .common-box {
    width: 100%;
    max-width: 800px;

    .Interation-box {
      display: flex;

      .main {
        width: 400px;
        padding: 30px 5px 30px 5px;
      }

      .button-box {
        top: 10px;
        display: flex;
        flex-direction: column;
        height: 300px;
        justify-content: space-around;
      }
    }
  }
}
</style>
